package server

import (
	"log"
	"sync"

	"github.com/aler9/gortsplib"
	"github.com/aler9/gortsplib/pkg/base"
)

// This example shows how to
// 1. create a RTSP server which accepts plain connections
// 2. allow a single client to publish a stream with TCP or UDP
// 3. allow multiple clients to read that stream with TCP, UDP or UDP-multicast

type RtspServerHandler struct {
	engine *WebRtcEngine

	sessionMutex sync.RWMutex
	sessionMap   map[string]*RtspSession
}

func NewRtspServerHandler(engine *WebRtcEngine) *RtspServerHandler {
	return &RtspServerHandler{
		engine:     engine,
		sessionMap: map[string]*RtspSession{},
	}
}

func (sh *RtspServerHandler) GetRtspConnection(name string) *RtspSession {
	sh.sessionMutex.Lock()
	defer sh.sessionMutex.Unlock()

	for _, conn := range sh.sessionMap {
		if conn.path == name {
			return conn
		}
	}

	return nil
}

func (sh *RtspServerHandler) GetRtspConnectionList() []*RtspSession {
	sh.sessionMutex.Lock()
	defer sh.sessionMutex.Unlock()

	var sessionList []*RtspSession
	for _, session := range sh.sessionMap {
		sessionList = append(sessionList, session)
	}

	return sessionList
}

// OnConnOpen called when a connection is opened.
func (sh *RtspServerHandler) OnConnOpen(ctx *gortsplib.ServerHandlerOnConnOpenCtx) {
	log.Printf("conn opened")
}

// OnConnClose called when a connection is closed.
func (sh *RtspServerHandler) OnConnClose(ctx *gortsplib.ServerHandlerOnConnCloseCtx) {
	log.Printf("conn closed (%v)", ctx.Error)
}

// OnSessionOpen called when a session is opened.
func (sh *RtspServerHandler) OnSessionOpen(ctx *gortsplib.ServerHandlerOnSessionOpenCtx) {
	log.Printf("session opened")
}

// OnSessionClose called when a session is closed.
func (sh *RtspServerHandler) OnSessionClose(ctx *gortsplib.ServerHandlerOnSessionCloseCtx) {
	log.Printf("session closed")

	sh.sessionMutex.Lock()
	defer sh.sessionMutex.Unlock()

	for _, session := range sh.sessionMap {
		if ctx.Session == session.publisher {
			session.Close()
			delete(sh.sessionMap, session.path)

			break
		}
	}
}

// OnDescribe called after receiving a DESCRIBE request.
func (sh *RtspServerHandler) OnDescribe(ctx *gortsplib.ServerHandlerOnDescribeCtx) (*base.Response, *gortsplib.ServerStream, error) {
	log.Printf("describe request")

	sh.sessionMutex.Lock()
	defer sh.sessionMutex.Unlock()

	session, ok := sh.sessionMap[ctx.Path]
	if !ok {
		log.Println("not found session ", ctx.Path)
		return &base.Response{
			StatusCode: base.StatusNotFound,
		}, nil, nil
	}

	// send the track list that is being published to the client
	return &base.Response{
		StatusCode: base.StatusOK,
	}, session.stream, nil
}

// OnAnnounce called after receiving an ANNOUNCE request.
func (sh *RtspServerHandler) OnAnnounce(ctx *gortsplib.ServerHandlerOnAnnounceCtx) (*base.Response, error) {
	log.Printf("announce request")

	sh.sessionMutex.Lock()
	defer sh.sessionMutex.Unlock()

	session := NewRtspSession(ctx)

	// save the track list and the publisher
	sh.sessionMap[session.path] = session

	return &base.Response{
		StatusCode: base.StatusOK,
	}, nil
}

// OnSetup called after receiving a SETUP request.
func (sh *RtspServerHandler) OnSetup(ctx *gortsplib.ServerHandlerOnSetupCtx) (*base.Response, *gortsplib.ServerStream, error) {
	log.Printf("setup request")

	session, ok := sh.sessionMap[ctx.Path]
	if !ok {
		log.Println("not found session ", ctx.Path)
		return &base.Response{
			StatusCode: base.StatusNotFound,
		}, nil, nil
	}

	return &base.Response{
		StatusCode: base.StatusOK,
	}, session.stream, nil
}

// OnPlay called after receiving a PLAY request.
func (sh *RtspServerHandler) OnPlay(ctx *gortsplib.ServerHandlerOnPlayCtx) (*base.Response, error) {
	log.Printf("play request")

	return &base.Response{
		StatusCode: base.StatusOK,
	}, nil
}

// OnRecord called after receiving a RECORD request.
func (sh *RtspServerHandler) OnRecord(ctx *gortsplib.ServerHandlerOnRecordCtx) (*base.Response, error) {
	log.Printf("record request")

	return &base.Response{
		StatusCode: base.StatusOK,
	}, nil
}

// OnPacketRTP called after receiving a RTP packet.
func (sh *RtspServerHandler) OnPacketRTP(ctx *gortsplib.ServerHandlerOnPacketRTPCtx) {
	sh.sessionMutex.RLock()
	defer sh.sessionMutex.RUnlock()

	// if we are the publisher, route the RTP packet to all readers
	for _, session := range sh.sessionMap {
		if ctx.Session == session.publisher {
			session.stream.WritePacketRTP(ctx.TrackID, ctx.Packet, ctx.PTSEqualsDTS)
			session.onPacketRTP(ctx) // 转给webrtc
			break
		}
	}
}
